package tagging;

import java.util.Vector;

public class ViterbiProcessor {
  public static final int MAX_SENTENCE_SIZE = 300;
  private static final double NNP_COEFFICIENT = 0.0015;

  TrainingSet ts_;
  Vector<String> tags_;
  int numTags_;

  public ViterbiProcessor(TrainingSet ts) {
    ts_ = ts;
    tags_ = ts.getTagStrings();
    numTags_ = tags_.size();
  }

  private boolean isNNP(String s) {
    return Character.isUpperCase(s.charAt(0));
  }

  private boolean isCD(String s) {
    try {
      Double.parseDouble(s);
    } catch (NumberFormatException e) {
      return false;
    }

    return true;
  }

  public void process(Vector<Parser.ProcessedWord> pwords) {
    boolean[][] valid = new boolean[MAX_SENTENCE_SIZE][numTags_];
    double[][] vvalue = new double[MAX_SENTENCE_SIZE][numTags_];
    int[][] prev = new int[MAX_SENTENCE_SIZE][numTags_];
    SuffixHandler sh = ts_.getSuffixHandler();

    for (int j = 0; j < numTags_; ++j) {
      if (tags_.get(j).equals(Tag.START_TAG)) {
        prev[0][j] = -1;
        valid[0][j] = true;
        vvalue[0][j] = 0.0;
      } else {
        valid[0][j] = false;
      }
    }


    for (int i = 1; i < pwords.size(); ++i) {
      String currWord = pwords.get(i).getWord();
      boolean isUnknown = true;
      for (int j = 0; j < numTags_; ++j) {
        String currTagString = tags_.get(j);
        Tag currTag = ts_.getTrainingTag(currTagString);

        valid[i][j] = false;
        
        if ((currTagString.equals("NNP") || currTagString.equals("NNPS")) &&
            !isNNP(pwords.get(i).getOriginalWord())) {
          continue;
        }

        double mx = 0.0;
        int pmax = -1;
        for (int p = 0; p < numTags_; ++p) {
          if (valid[i-1][p]) {
            Tag prevTag = ts_.getTrainingTag(tags_.get(p));
            int bigramCount = prevTag.getTagCount(currTagString);

            if (bigramCount == 0) continue;

            double num = Math.log((double)bigramCount) + vvalue[i-1][p];
            double den = Math.log((double)prevTag.getTotalCount());
            double val = num - den;

            if (pmax == -1 || mx < val) {
              mx = val;
              pmax = p;
            }
          }
        }

        prev[i][j] = pmax;
        if (pmax >= 0) {
          valid[i][j] = true;
          int wordCount = currTag.getWordCount(currWord);

          if (wordCount >= TrainingSet.UNKNOWN_THRESHOLD) {
            isUnknown = false;
          }  

          if (isNNP(pwords.get(i).getOriginalWord()) && i > 1 && currTagString.equals("NNP")) {
            // Capitalization in the middle of the sentence.
            vvalue[i][j] = mx + Math.log(NNP_COEFFICIENT);
          } else if (isCD(currWord) && currTagString.equals("CD")) {
            // The whole word is a number.
            vvalue[i][j] = mx;
          } else {
            if (wordCount >= TrainingSet.UNKNOWN_THRESHOLD) {
              vvalue[i][j] = mx + Math.log((double)wordCount) - Math.log((double)currTag.getTotalCount());
            } else {
              /*
              vvalue[i][j] = mx
              + Math.log((double)currTag.getGTWordCount(wordCount)) - Math.log((double)currTag.getTotalCount());
              */
              
              int suffixCount = sh.getSuffixCount(currWord, currTagString);
              
              if (suffixCount > 0) {
                vvalue[i][j] = mx
                + Math.log((double)currTag.getGTWordCount(wordCount)) - Math.log((double)currTag.getTotalCount())
                + Math.log((double) suffixCount) - Math.log((double)currTag.getTotalCount());
              } else {
                vvalue[i][j] = mx
                + Math.log((double)currTag.getGTWordCount(wordCount)) - Math.log((double)currTag.getTotalCount())
                + Math.log((double) 0.1) - Math.log((double)currTag.getTotalCount());
              }
            }
          }
        } else {
          valid[i][j] = false;
        }
      }

      if (isUnknown) {
        pwords.get(i).setUnknownBit();
      }
    }

    int jmax = -1;
    double mxv = 0.0;
    int lastIndex = pwords.size() - 1;
    for (int j = 0; j < numTags_; ++j) {
      if (valid[lastIndex][j]) {
        if (jmax == -1 || mxv < vvalue[lastIndex][j]) {
          mxv = vvalue[lastIndex][j];
          jmax = j;
        }
      }
    }

    int tagIndex = jmax;
    for (int i = lastIndex; i >= 0; --i) {
      pwords.get(i).setTag(tags_.get(tagIndex));
      tagIndex = prev[i][tagIndex];
    }
  }
}
